#include <iostream>
#include <fstream>
#include <random>
#include <unordered_set>
#include "hnsw.hpp"
#include "vecs_io.hpp"
#include "ground_truth.hpp"
#include "time_memory.hpp"

using namespace std;
using namespace MultipleHNSW;


int main(int argc, char **argv) {
    if (argc != 2) {
        std::cout << argv[0] << " base_fname" << std::endl;
        exit(-1);
    }
    printf("base_fname %s", argv[1]);

    char *dataset_name = argv[1];
    const int M = 8;
    const int ef_construction = 100;
    const int random_seed = 100;
    const int n_index = 5;

    const int ef_search_config[3] = {10, 5, 50}; //start, stride, iter

    int n_query, top_k, n_base, data_dim;
    top_k = 10;

    vector<float *> data = read_data(dataset_name, n_query, top_k, n_base, data_dim);
    float *base = data[0];
    float *query = data[1];
    int *gnd = (int *) data[2];

    printf("M %d, efConstruction %d, n_index %d\n", M, ef_construction, n_index);

    //init hnsw
    vector<int> indexing_time_l(n_index);
    HNSW *hnsw_l = new HNSW[n_index];
    // create shuffle array
    vector<vector<int>> idx_shuffle_l_l(n_index);

    for (int i = 0; i < n_index; i++) {
        idx_shuffle_l_l[i] = vector<int>(n_base);
        if (i == 0) {
            for (int j = 0; j < n_base; j++) {
                idx_shuffle_l_l[i][j] = j;
            }
        } else {
            for (int j = 0; j < n_base; j++) {
                idx_shuffle_l_l[i][j] = idx_shuffle_l_l[i - 1][j];
            }
            std::shuffle(idx_shuffle_l_l[i].begin(), idx_shuffle_l_l[i].end(),
                         default_random_engine(random_seed + i));
        }

    }

    for (int i = 0; i < n_index; i++) {
        TimeRecord total_insert_record;
        HNSW hnsw = HNSW(data_dim, n_base, M, ef_construction, random_seed + i);
        printf("build index %d\n", i);
        hnsw.insert(base, idx_shuffle_l_l[i]);
        //hnsw.test_graph();
        indexing_time_l[i] = (int) (total_insert_record.get_elapsed_time_micro() * 1e-6);
        total_insert_record.reset();
        hnsw_l[i] = hnsw;
    }

    vector<pair<double, double>> time_recall_l(ef_search_config[2]);
    vector<int> ef_search_l(ef_search_config[2]);

    printf("querying\n");

    for (int iter = 0, ef_search = ef_search_config[0];
         iter < ef_search_config[2]; iter++, ef_search += ef_search_config[1]) {

        ef_search_l[iter] = ef_search;
        double query_time = 0;

        vector<unordered_set<int>> test_gnd_l(n_query);

        for (int idx = 0; idx < n_index; idx++) {
            hnsw_l[idx].ef_search_ = ef_search;
            TimeRecord query_record;
            vector<unordered_set<int>> tmp_test_gnd_l = hnsw_l[idx].search_knn(query, top_k, n_query);

            double tmp_query_time = query_record.get_elapsed_time_micro() / n_query * 1e-3; //unit ms
            query_time += tmp_query_time;

            for (int i = 0; i < n_query; i++) {
                test_gnd_l[i].insert(tmp_test_gnd_l[i].begin(), tmp_test_gnd_l[i].end());
            }

        }

        double recall = count_recall(n_query, top_k, test_gnd_l, gnd);
        printf("efSearch %d, average recall: %.3f, query time %.1f ms\n", ef_search, recall, query_time);
        time_recall_l[iter] = make_pair(query_time, recall);
    }

    size_t peak_memory = get_peak_RSS();
    printf("peak memory %ld Mb\n", peak_memory / 1000000);

    SaveNode save_node;
    save_node.dataset = dataset_name;
    save_node.method_name = "cascade_shuffle";
    save_node.top_k = top_k;
    save_node.time_recall_l = time_recall_l;
    save_node.ef_search_l = ef_search_l;
    save_node.n_index = n_index;
    save_node.total_M = M * n_index;
    save_node.ef_construction = ef_construction;
    save_node.indexing_time_l = indexing_time_l;
    save_node.max_memory = peak_memory / 1000000;

    vector<int> M_l;
    for (int i = 0; i < n_index; i++) {
        M_l.push_back(M);
    }
    save_node.M_l = M_l;
    save_node.save();

    delete[] hnsw_l;
    return 0;
}